import os, glob, string, csv
import time, datetime
import pytz
import numpy as np

from pytz import timezone
from scipy import ndimage, misc
from numpy import *

# make this configurable.
img_dir  = './images'
img_pfix = '*.jpg'
tz_pst   = timezone('US/Pacific')
time_fmt = '%Y%m%d%H%M'
img_cols = ['time', 'irbw', 'vis', 'icg', 'nir', 'wv']
wea_cols = ['time', 'wind_dir', 'wind_speed', 'temp', 'dew', 'pressure']
irr_cols = ['time', 'ghi', 'dni', 'diff', 'temp', 'ws', 'zs', 'azs', 'am']

def main():
    # make below parameters as command-line args
    weather_src       = './csv/weatherData3.csv'
    irrad_src         = './csv/irrad.txt'
    img_data_file     = './data/image_data.csv'
    weather_data_file = './data/weather_data.csv'
    irrad_data_file   = './data/irrad.csv'
    out_data          = './data/output.csv'

    start, stop = generateImgData(img_dir, img_pfix, img_data_file)
    generateWeatherData(start, stop, weather_src, weather_data_file)
    generateIrradData(start, stop, irrad_src, irrad_data_file)
    mergeData(img_data_file, weather_data_file, irrad_data_file, out_data)

def generateIrradData(start, stop, irrad_src, irrad_data_file):
    with open(irrad_src, 'r') as src:
        out = open(irrad_data_file, 'w')
        writer = csv.writer(out, delimiter=',')
        tfmt =  '%m/%d/%Y %H:%M'

        # read header and write it as csv
        line = src.readline()
        writer.writerow(irr_cols)

        while True:
            line = src.readline()
            if not line: break
            token = line.split(',')

            # check time
            t = time.strptime(token[0] + ' ' + token[1], tfmt)
            t = datetime.datetime(*t[:6]).replace(tzinfo=tz_pst)

            if t < start: continue
            elif t > stop: break

            token[-1] = float(token[-1])
            writer.writerow([t.strftime(time_fmt)] + token[2:])

def generateImgData(_path, _postfix, _fname):
    # list image files
    path  = os.path.join(_path, img_pfix)
    flist = glob.glob(path);
    flist.sort()

    # setup csv output
    data_out  = open(_fname, 'w')
    datawriter = csv.writer(data_out, delimiter=',')
    datawriter.writerow(img_cols)

    # process images depend on its name.
    prev_time = time.strptime('190001010000', time_fmt)
    count = 0
    total = 0
    written = 0
    collect = True
    IsFirst = False
    start_time = None
    end_time = None

    # load csv files, TODO: make these as settable.
    nir_idx = np.genfromtxt(open('csv/nir_index.csv', 'r'),     \
                            delimiter=',', dtype='f8')
    icg_idx = np.genfromtxt(open('csv/icg_index.csv', 'r'),     \
                            delimiter=',', dtype='f8')
    wv_idx  = np.genfromtxt(open('csv/wv_index.csv', 'r'),      \
                            delimiter=',', dtype='f8')

    # initiate the time
    token = string.split(flist[0], '/')[2]
    token = string.split(token, '_')
    prev_time = time.strptime(token[0] + token[1], time_fmt) 
    values = [0, 0, 0, 0, 0]
    for f in flist:
        # parse a file name (date, time, region, type)
        _fname = string.split(f, '/')[2]
        _token = string.split(_fname, '_')

        # create time object
        _time = time.strptime(_token[0] + _token[1], time_fmt) 
        if _time != prev_time:
            if count == 5:
                # handling timezone
                ct = datetime.datetime(*prev_time[:6]).replace(tzinfo=pytz.utc)
                if ct.hour == 0 or ct.hour == 1 or ct.hour == 2:
                    # need to adjust day
                    ct = ct - datetime.timedelta(hours=ct.hour + 1)

                # write file
                datawriter.writerow([ct.astimezone(tz_pst).strftime(time_fmt),  \
                                    values[0], values[1], values[2], values[3]])
                if IsFirst != True:
                    start_time = ct.astimezone(tz_pst)
                    IsFirst = True
                else:
                    end_time = ct.astimezone(tz_pst)

                written = written + 5
            count = 0
            values = [0, 0, 0, 0, 0]
            prev_time = _time

        # region
        _region = _token[2]

        # image type
        _type = string.split(_token[3], '.')[0]

        # read image
        if _region == 'US':
            _img = misc.imread(f)[506:509, 291:294]
        elif _region == 'LAS':
            _img = misc.imread(f)[349:352, 312:315]
        else:
            print 'Unknown region: %s, skipped' % _region

        # calculate a value
        if _type == 'icg':
            v = findBestMatch(_img, icg_idx)
        elif _type == 'nir':
            v = findBestMatch(_img, nir_idx)
        elif _type == 'wv':
            v = findBestMatch(_img, wv_idx)
        else:
            # just average pixels
            v = _img.mean()

        values[count] = v
        count = count + 1
        total = total + 1

    return start_time, end_time

def generateWeatherData(_start, _stop, _src, _dst):
    # start processing weather data
    weather = csv.reader(open(_src, 'rU'), delimiter=',')

    # take a data only in the range
    weather_out  = open(_dst, 'w')
    weatherwriter = csv.writer(weather_out, delimiter=',')
    weatherwriter.writerow(wea_cols)
    interest_idx = array([2, 3, 4, 6, 21, 22, 25])
    for line in weather:
        try:
            aTime = time.strptime(line[2], time_fmt)
            aTime = datetime.datetime(*aTime[:6]).replace(tzinfo=tz_pst)
        except:
            # ignore exceptions
            continue

        if aTime < _start:
            continue
        elif aTime > _stop: break

        # valid timestamp
        elem = np.array(line)[interest_idx]
        weatherwriter.writerow(elem)

def mergeData(_fname_img, _fname_weather, _fname_irrad, _fname_out):
    _out  = open(_fname_out, 'w')
    _writer = csv.writer(_out, delimiter=',')
    _writer.writerow(['label'] + img_cols + wea_cols[1:] + irr_cols[1:])
    last_good = None

    with open(_fname_img, 'r') as img_data:
        # for weather data
        img_reader = csv.reader(img_data, delimiter=',')
        weather_reader = csv.reader(open(_fname_weather, 'r'), delimiter=',')

        # for image data
        line = img_reader.next()
        line = img_reader.next()
        img_time = time.strptime(line[0], time_fmt)
        img_time = datetime.datetime(*img_time[:6])

        # for irrad data
        irrad_reader = csv.reader(open(_fname_irrad, 'r'), delimiter=',')
        irrad = irrad_reader.next()
        irrad = irrad_reader.next()
        irrad_time = time.strptime(irrad[0], time_fmt)
        irrad_time = datetime.datetime(*irrad_time[:6])

        for entry in weather_reader:
            # get a time
            try:
                weather_time = time.strptime(entry[0], time_fmt)
                weather_time = datetime.datetime(*weather_time[:4])
            except:
                continue

            while True:
                try:
                    # round up img_time
                    r_time = img_time + datetime.timedelta(minutes=30)
                    r_time = r_time - datetime.timedelta(
                                                minutes=r_time.minute,
                                                seconds=r_time.second,
                                                microseconds=r_time.microsecond
                                                )
                    if weather_time > r_time:
                        line =  img_reader.next()
                        img_time = time.strptime(line[0], time_fmt)
                        img_time = datetime.datetime(*img_time[:6])
                    else:
                        break
                except StopIteration:
                    break
                except:
                    continue

            if weather_time != r_time: continue

            # find a match
            if validateData(entry):
                datum = line + entry[1:]
                last_good = entry[1:]
            else:
                # data has invalidate field, recover
                if not last_good == None:
                    # use last good record
                    datum = line + recoverData(entry[1:], last_good)
                else:
                    print 'skip: ', entry
                    continue

             # search irrad data
            while True:
                try:
                    if img_time > irrad_time:
                        irrad = irrad_reader.next()
                        irrad_time = time.strptime(irrad[0], time_fmt)
                        irrad_time = datetime.datetime(*irrad_time[:6])
                    else:
                        break
                except StopIteration:
                    break
                except:
                    continue

            # calculate a label
            GHI, CF = estimateLabel(irrad[1], irrad[6])
            if CF > 0.6:
                label = 1
            else:
                label = 0

            _writer.writerow([label] + datum + irrad[1:])

def validateData(data):
    for field in data:
        if not field.find('*') < 0: return False
    return True

def recoverData(data, good):
    for idx in range(0, len(data)):
        if not data[idx].find('*') < 0:
            data[idx] = good[idx]
    return data

def estimateLabel(GHI, Zs):
    import math

    # TODO: check whether or not we need below if-statement.
    #if Zs > 89:
    #    return [0, 0]

    logGHI = 7.11179 + (1.031893 * math.log(math.cos(math.radians(float(Zs))))) \
           - (0.04543362 / math.cos(math.radians(float(Zs))))
    GHI_clear = math.exp(logGHI)
    CF = 1 - (float(GHI) / GHI_clear)

    return GHI_clear, CF

def findBestMatch(_v, _idx):
    if len(_v.shape) != 2 and len(_v.shape) != 3:
        print 'Input is in wrong shape, ', _v.shape
        return -999999 # invalid data

    # calculate a norm
    _avg = np.average(np.average(_v, 1), 0)
    _dist = _idx[:, 1:] - np.tile(_avg, [_idx.shape[0], 1])
    _dist = np.sum(np.sqrt(np.power(_dist, 2)), axis=1)

    return _idx[np.argmin(_dist), 0]

if __name__ == '__main__':
    main()
